package top.longmarch.flow.ai;

import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.catalina.connector.ClientAbortException;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.Serializable;
import java.nio.charset.StandardCharsets;
import java.time.Duration;

/**
 * 处理长时间运行的异步 HTTP 请求工具类
 * <p>
 * 包括 {@link ResponseBodyEmitter} 和 {@link SseEmitter}，
 * 其中 SSE（Server-Send Event）服务端发送事件，用于大模型对话流式应答
 */
@Slf4j
public class AsyncResponseUtil {

    /**
     * 默认超时时间
     */
    public static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(120);
    /**
     * {@link MediaType} text/event-stream;charset=UTF-8
     */
    public static final MediaType TEXT_EVENT_STREAM_UTF8 = new MediaType(MediaType.TEXT_EVENT_STREAM, StandardCharsets.UTF_8);
    /**
     * text/event-stream;charset=UTF-8
     */
    public static final String TEXT_EVENT_STREAM_UTF8_VALUE = TEXT_EVENT_STREAM_UTF8.toString();
    /**
     * 默认 SSE 异常事件 name
     */
    public static final String ERROR_EVENT = "error";

    /**
     * 创建 Response 发送器，并设置默认超时时间，超时自动关闭连接
     *
     * @return ResponseBodyEmitter
     */
    public static ResponseBodyEmitter createResponseBodyEmitter() {
        return createResponseBodyEmitter(DEFAULT_TIMEOUT, null);
    }

    /**
     * 创建 Response 发送器
     *
     * @param timeout         超时时间
     * @param timeoutCallback 超时回调，为 null 时调用 {@link ResponseBodyEmitter#complete()} 方法
     * @return ResponseBodyEmitter
     */
    public static ResponseBodyEmitter createResponseBodyEmitter(Duration timeout, Runnable timeoutCallback) {
        ResponseBodyEmitter emitter = new ResponseBodyEmitter(timeout.toMillis());
        if (timeoutCallback == null) {
            emitter.onTimeout(emitter::complete);
        } else {
            emitter.onTimeout(timeoutCallback);
        }
        return emitter;
    }

    /**
     * 创建 UTF-8 编码的 SSE 发送器，以避免默认编码出现中文乱码，并设置默认超时时间，超时自动关闭连接
     *
     * @return SseEmitter
     */
    public static SseEmitter createUtf8SseEmitter() {
        return createUtf8SseEmitter(DEFAULT_TIMEOUT, null);
    }

    /**
     * 创建 UTF-8 编码的 SSE 发送器，以避免默认编码出现中文乱码
     *
     * @param timeout         超时时间
     * @param timeoutCallback 超时回调，为 null 时调用 {@link SseEmitter#complete()} 方法
     * @return SseEmitter
     */
    public static SseEmitter createUtf8SseEmitter(Duration timeout, Runnable timeoutCallback) {
        SseEmitter emitter = new SseEmitter(timeout.toMillis()) {
            @Override
            protected void extendResponse(ServerHttpResponse outputMessage) {
                super.extendResponse(outputMessage);
                HttpHeaders headers = outputMessage.getHeaders();
                headers.setContentType(TEXT_EVENT_STREAM_UTF8);
            }
        };
        if (timeoutCallback == null) {
            emitter.onTimeout(emitter::complete);
        } else {
            emitter.onTimeout(timeoutCallback);
        }
        return emitter;
    }

    /**
     * 发送 SSE 数据
     *
     * @param emitter SseEmitter
     * @param data    SSE 协议参数-data，对象默认转换为 JSON 文本
     * @return 成功-true，失败-false
     */
    public static boolean send(SseEmitter emitter, Object data) {
        return send(emitter, null, null, data);
    }

    /**
     * 发送 SSE 数据
     *
     * @param emitter SseEmitter
     * @param id      SSE 协议参数-id
     * @param data    SSE 协议参数-data，对象默认转换为 JSON 文本
     * @return 成功-true，失败-false
     */
    public static boolean send(SseEmitter emitter, Serializable id, Object data) {
        return send(emitter, id, null, data);
    }

    /**
     * 发送 SSE 数据
     *
     * @param emitter SseEmitter
     * @param id      SSE 协议参数-id
     * @param name    SSE 协议参数-name
     * @param data    SSE 协议参数-data，对象默认转换为JSON文本
     * @return 成功-true，失败-false
     */
    public static boolean send(SseEmitter emitter, Serializable id, String name, Object data) {
        try {
            SseEmitter.SseEventBuilder event = SseEmitter.event();
            if (id != null) {
                event.id(String.valueOf(id));
            }
            if (name != null) {
                event.name(name);
            }
            if (data != null) {
                if (data instanceof CharSequence) {
                    event.data(((CharSequence) data).toString());
                } else {
                    event.data(JSONUtil.toJsonStr(data));
                }
            }
            emitter.send(event);
            return true;
        } catch (ClientAbortException e) {
            log.error("客户端主动关闭连接: {}", e.getMessage());
            return false;
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            return false;
        }
    }

    /**
     * 发送异步数据
     *
     * @param emitter ResponseBodyEmitter
     * @param data    数据
     * @return 成功-true，失败-false
     */
    public static boolean send(ResponseBodyEmitter emitter, Object data) {
        return send(emitter, data, null);
    }

    /**
     * 发送异步数据
     *
     * @param emitter   ResponseBodyEmitter
     * @param data      数据
     * @param mediaType 响应类型，自动选择合适的 HttpMessageConverter
     * @return 成功-true，失败-false
     */
    public static boolean send(ResponseBodyEmitter emitter, Object data, MediaType mediaType) {
        try {
            emitter.send(data, mediaType);
            return true;
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            return false;
        }
    }

    /**
     * 关闭异步发送
     *
     * @param emitter 响应发送器
     */
    public static void complete(ResponseBodyEmitter emitter) {
        emitter.complete();
    }

    /**
     * 转换为 WrappedResponseBodyEmitter
     *
     * @param emitter 响应发送器
     * @return WrappedResponseBodyEmitter
     */
    public static <T extends ResponseBodyEmitter> WrappedResponseBodyEmitter wrapEmitter(T emitter) {
        return WrappedResponseBodyEmitter.wrap(emitter);
    }

    /**
     * 封装 ResponseBodyEmitter
     */
    public static class WrappedResponseBodyEmitter {

        /**
         * SseEmitter
         */
        private ResponseBodyEmitter emitter;

        /**
         * 连接是否关闭
         */
        private boolean completed;

        public static <T extends ResponseBodyEmitter> WrappedResponseBodyEmitter wrap(T emitter) {
            WrappedResponseBodyEmitter wrapped = new WrappedResponseBodyEmitter();
            wrapped.emitter = emitter;
            wrapped.emitter.onCompletion(wrapped::complete);
            return wrapped;
        }

        public <T extends ResponseBodyEmitter> T unwrap() {
            return (T) this.emitter;
        }

        public void onComplete(Runnable callback) {
            emitter.onCompletion(callback);
        }

        public void complete() {
            this.completed = true;
        }

        public boolean isCompleted() {
            return this.completed;
        }

    }

}
